package edu.buaa.agent;

import edu.buaa.agent.intecepter.MySystemInterceptor;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.dynamic.loading.ClassReloadingStrategy;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.matcher.StringMatcher;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.nio.file.Files;

import static java.util.Collections.singletonMap;
import static net.bytebuddy.dynamic.ClassFileLocator.ForClassLoader.read;
import static net.bytebuddy.dynamic.loading.ClassInjector.UsingInstrumentation.Target.BOOTSTRAP;
import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * https://segmentfault.com/a/1190000007253689
 * https://www.infoq.cn/article/Easily-Create-Java-Agents-with-ByteBuddy/
 * https://blog.csdn.net/undergrowth/article/details/86493336
 * https://www.jianshu.com/p/7a5a2c78dab4
 * https://blog.csdn.net/DorMOUSENone/article/details/81781131
 * jvmAop 处理
 *
 * @author 汪小哥
 * @date 01-02-2020
 */
public class JvmAgentAop {
    private static final Class<?> INTERCEPTOR_CLASS = MySystemInterceptor.class;


    public static void main(String[] args) throws IOException {

        // 1. 安装 Java Agent
        Instrumentation instrumentation = ByteBuddyAgent.install();

        // 2. 对 System 类中的方法进行拦截增强
        new ByteBuddy()
                .redefine(System.class)
                .visit(Advice.to(CurrentTimeInterceptor.class).on(ElementMatchers.named("currentTimeMillis")))
                .make()
                .load(System.class.getClassLoader(),
                        ClassReloadingStrategy.of(instrumentation).reset(System.class));

        // 测试拦截效果
        System.out.println(System.currentTimeMillis());
    }

    public static class CurrentTimeInterceptor {

        // 在方法执行前增强逻辑
        @Advice.OnMethodEnter
        public static void onEnter() {
            System.out.println("方法调用前：准备获取时间戳");
        }

        // 在方法执行后增强逻辑
        @Advice.OnMethodExit(onThrowable = Throwable.class)
        public static void onExit(@Advice.Return(readOnly = false) long returnValue) {
            System.out.println("方法调用后：原始时间戳为 " + returnValue);
            returnValue += 1000; // 示例：修改返回值，增加 1 秒
            System.out.println("方法调用后：修改后时间戳为 " + returnValue);
        }

    }

//    public static void premain(String agentArgs, Instrumentation inst) throws IOException {
//        System.out.println("this is an agent.");
//        System.out.println("this is an agent 和jvm 一起启动");
//        System.out.println("args:" + agentArgs + "\n");
//        // 2. 对 System 类中的方法进行拦截增强
////        ByteBuddy byteBuddy = new ByteBuddy();
////                byteBuddy
////                .redefine(System.class)
////                .visit(Advice.to(CurrentTimeInterceptor.class).on(named("currentTimeMillis")))
////                .make()
////                .load(System.class.getClassLoader(),
////                        ClassReloadingStrategy.of(inst).reset(System.class));
////        new AgentBuilder.Default().with(byteBuddy)
//////                .redefine(System.class)
////                .(System.class)
////                .transform(new AgentBuilder.Transformer() {
////                    @Override
////                    public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
////                        return builder.method(ElementMatchers.any())
////                                .intercept(MethodDelegation.to(CurrentTimeInterceptor.class));
////                    }
////                })
////            .installOn(inst);
//
//        new AgentBuilder.Default()
//                .type(nameStartsWith("java.lang").and(named("System")))
//                .transform(new AgentBuilder.Transformer() {
//                    @Override
//                    public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
//                        return builder.method(ElementMatchers.any())
//                                .intercept(MethodDelegation.to(CostTimeInterceptor.class));
//                    }
//                }).installOn(inst);
//        System.out.println(System.currentTimeMillis());
//    }

    // From https://stackoverflow.com/questions/44747219/byte-buddy-advice-onmethodexit-constructor-retransformation
    public static void premain(String arg, Instrumentation instrumentation) throws Exception {
        System.out.println("this is an agent.");
        System.out.println("this is an agent 和jvm 一起启动");
        System.out.println("args:" + arg + "\n");
        injectBootstrapClasses(instrumentation);
        new AgentBuilder.Default()

//                .enableNativeMethodPrefix("foo")
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
                .ignore(new AgentBuilder.RawMatcher.ForElementMatchers(nameStartsWith("net.bytebuddy.").or(isSynthetic()), any(), any()))
                .with(new AgentBuilder.Listener.Filtering(
                        new StringMatcher("java.lang.System", StringMatcher.Mode.EQUALS_FULLY),
                        AgentBuilder.Listener.StreamWriting.toSystemOut()))
                .type(named("java.lang.System"))
                .transform((builder, type, classLoader, module,protectionDomain) ->
                        builder.visit(Advice.to(INTERCEPTOR_CLASS).on(named("currentTimeMillis")))
                )
                .installOn(instrumentation);
        // 测试拦截效果
        System.out.println(System.currentTimeMillis());
    }

    private static void injectBootstrapClasses(Instrumentation instrumentation) throws IOException {
        File temp = Files.createTempDirectory("tmp").toFile();
        temp.deleteOnExit();

        ClassInjector.UsingInstrumentation.of(temp, BOOTSTRAP, instrumentation)
                .inject(singletonMap(new TypeDescription.ForLoadedType(INTERCEPTOR_CLASS), read(INTERCEPTOR_CLASS)));
    }
}
